package top.v5it.japi.ddd.core.web;

import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.ttl.TransmittableThreadLocal;
import org.slf4j.MDC;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * 链路追踪
 *
 * @author zhanpu
 * @date 2022/8/26
 */
public final class TraceContextHolder {

    public static final String HEAD_TRACE_ID_NAME = "X-Trace-ID";
    public static final String HEAD_SPAN_ID_NAME = "X-Span-ID";
    static final String SLF4J_TRACE_ID_NAME = "X-TraceId";
    static final String SLF4J_SPAN_ID_NAME = "X-SpanId";
    static final String SLF4J_IP_NAME = "ip";
    static final String ROOT_SPAN_ID = "0";
    static final String SERVER_NAME = "server.name";

    private static final ThreadLocal<Long> REQUEST_TIME_ATTRIBUTES_HOLDER = new TransmittableThreadLocal<>(Boolean.TRUE);

    /**
     * 开始计时
     */
    static void start() {
        final long startTime = System.nanoTime();
        REQUEST_TIME_ATTRIBUTES_HOLDER.set(startTime);
    }

    /**
     * 计时结束
     *
     * @return
     */
    static long stopAndGet() {
        final long endTime = System.nanoTime();
        final Long startTime = REQUEST_TIME_ATTRIBUTES_HOLDER.get();
        Assert.notNull(startTime, () -> new IllegalArgumentException("必须先调用start方法"));
        REQUEST_TIME_ATTRIBUTES_HOLDER.remove();
        final long costTime = endTime - startTime;
        return TimeUnit.MILLISECONDS.convert(costTime, TimeUnit.NANOSECONDS);
    }

    /**
     * 获取ip
     *
     * @return
     */
    static String getIp() {
        return MDC.get(SLF4J_IP_NAME);
    }

    /**
     * 获取traceId
     *
     * @return
     */
    public static String getTraceId() {
        return MDC.get(SLF4J_TRACE_ID_NAME);
    }

    /**
     * 获取spanId
     *
     * @return
     */
    public static String getSpanId() {
        return MDC.get(SLF4J_SPAN_ID_NAME);
    }

    /**
     * 添加traceId
     *
     * @param traceId
     */
    static void putTraceId(String traceId) {
        MDC.put(SLF4J_TRACE_ID_NAME, traceId);
    }

    /**
     * 添加spanId
     *
     * @param spanId
     */
    static void putSpanId(String spanId) {
        MDC.put(SLF4J_SPAN_ID_NAME, spanId);
    }

    /**
     * 添加ip
     *
     * @param ip
     */
    static void putIp(String ip) {
        MDC.put(SLF4J_IP_NAME, ip);
    }

    /**
     * 删除traceId
     */
    static void removeTraceId() {
        MDC.remove(SLF4J_TRACE_ID_NAME);
    }

    /**
     * 删除ip
     */
    static void removeIp() {
        MDC.remove(SLF4J_IP_NAME);
    }

    /**
     * 删除spanId
     */
    static void removeSpanId() {
        MDC.remove(SLF4J_SPAN_ID_NAME);
    }

    /**
     * 计算并设置traceId
     *
     * @return
     */
    public static String computeTraceId() {
        return Optional
                .ofNullable(getTraceId())
                .filter(StrUtil::isNotBlank)
                .orElse(IdUtil.randomUUID());
    }

    /**
     * 生成spanId
     *
     * @param spanId
     * @return
     */
    static String genSpanId(String spanId) {
        return StrUtil.isBlank(spanId) ? ROOT_SPAN_ID : spanId + ".1";
    }

    /**
     * 获取根spanId
     *
     * @return
     */
    public static String rootSpanId() {
        return ROOT_SPAN_ID;
    }
}
